import os

from dotenv import load_dotenv
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI

# 加载.env文件
load_dotenv()

# 调用大语言模型
# 创建模型
model = ChatOpenAI(
	model='qwen-plus',
	api_key=os.getenv("OPENAI_API_KEY"),
	base_url=os.getenv("BASE_URL"),
)

# 获取解析器
parser = StrOutputParser()

# 定义提示模板
prompt_template = ChatPromptTemplate.from_messages([
	('system', '你是一个乐于助人的助手。用{language}尽你所能回答所有问题。'),
	MessagesPlaceholder(variable_name='my_msg')
])

# 获得链
chain = prompt_template | model | parser

# 保存聊天的历史记录
store = {}  # 所有用户的聊天记录都保存到store。key: sessionId,value: 历史聊天记录对象store = {}  # 所有用户的聊天记录都保存到store。key: sessionId,value: 历史聊天记录对象


# 此函数预期将接收一个session_id并返回一个消息历史记录对象。
def get_session_history(session_id: str):
	if session_id not in store:
		store[session_id] = ChatMessageHistory()
	return store[session_id]


do_message = RunnableWithMessageHistory(chain, get_session_history, input_messages_key='my_msg')

config = {'configurable': {'session_id': 'zs1234'}}  # 给当前会话定义一个sessionId

# 第一轮
resp1 = do_message.invoke(
	{
		'my_msg': [HumanMessage(content='你好啊！ 我是LaoXiao')],
		'language': '中文'
	},
	config=config
)

print(resp1)

# 第二轮
resp2 = do_message.invoke(
	{
		'my_msg': [HumanMessage(content='请问：我的名字是什么？')],
		'language': '中文'
	},
	config=config
)

print(resp2)

# 第3轮： 返回的数据是流式的
config = {'configurable': {'session_id': 'lis2323'}}  # 给当前会话定义一个sessionId
for resp in do_message.stream({'my_msg': [HumanMessage(content='请给我讲一个笑话？')], 'language': '中文'},
                              config=config):
	# 每一次resp都是一个token
	print(resp, end='-')
